import socket
import sqlite3
from werkzeug.security import generate_password_hash
from werkzeug.security import check_password_hash


class UserDataBase():
    def __init__(self):
        self.conn = sqlite3.connect('.\\User.db')
        self.cursor = self.conn.cursor()
    def add_user(self, username, password):#注册提交，新建用户
        self.cursor.execute('SELECT * FROM users WHERE username = ?', (username,))
        user = self.cursor.fetchone()
        if user is not None:#用户名已存在，返回False
            return False
        self.cursor.execute('INSERT INTO users (username, password) VALUES (?, ?)', (username, generate_password_hash(password)))
        self.conn.commit()
        return True#完成注册返回True
    def check_login(self,username, password):#登录验证
        self.cursor.execute('SELECT * FROM users WHERE username = ?', (username,))
        user = self.cursor.fetchone()
        if user is None:
            return False
        stored_password = user[2]
        if user[0] and check_password_hash(stored_password, password):
            return True
        return False
    def check_client(self,client_socket):#客户端验证

        get_message = client_socket.recv(1024).decode()
        type,username,password = get_message.split(" ")
        if type == "login":#登录操作
            if self.check_login(username, password):
                client_socket.sendall(b"Authentication successful!")
                print(f"login,{username} Logged in")
            else:
                print(f"loginFailed")
                client_socket.sendall(b"Authentication failed.")
        elif type == "register":#注册操作
            if self.add_user(username, password):
                client_socket.sendall(b"Registration successful!")
                print(f"register,{username} Registered")
            else:
                client_socket.sendall(b"Registration failed.")
                print(f"registerFailed")
        client_socket.close()

if __name__ == "__main__":
    # 创建TCP套接字
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.bind(("0.0.0.0", 23621))  # 绑定到所有网络接口上的指定端口
    server_socket.listen(5)  # 开始监听，允许最多5个连接
    DB = UserDataBase()
    print("等待客户端连接...")

    while True:
        # 接受客户端连接
        client_socket, addr = server_socket.accept()
        print("接受来自 %s 的连接" % str(addr))
        DB.check_client(client_socket)